import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras import layers
import os
import glob
import time
import warnings
warnings.filterwarnings('ignore')  # 忽略所有警告
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False    # 正确显示负号
# 设置随机种子
tf.random.set_seed(42)
np.random.seed(42)

# 超参数
BUFFER_SIZE = 1000
BATCH_SIZE = 1
IMG_WIDTH = 256
IMG_HEIGHT = 256
INPUT_CHANNELS = 3
OUTPUT_CHANNELS = 3
LAMBDA = 10
EPOCHS = 50


# 生成器模型 (U-Net架构)
def build_generator():
    inputs = tf.keras.layers.Input(shape=[IMG_HEIGHT, IMG_WIDTH, INPUT_CHANNELS])

    # 下采样
    down1 = layers.Conv2D(64, 4, strides=2, padding='same', activation='relu')(inputs)
    down2 = layers.Conv2D(128, 4, strides=2, padding='same', activation='relu')(down1)
    down2 = layers.BatchNormalization()(down2)
    down3 = layers.Conv2D(256, 4, strides=2, padding='same', activation='relu')(down2)
    down3 = layers.BatchNormalization()(down3)
    down4 = layers.Conv2D(512, 4, strides=2, padding='same', activation='relu')(down3)
    down4 = layers.BatchNormalization()(down4)
    down5 = layers.Conv2D(512, 4, strides=2, padding='same', activation='relu')(down4)
    down5 = layers.BatchNormalization()(down5)
    down6 = layers.Conv2D(512, 4, strides=2, padding='same', activation='relu')(down5)
    down6 = layers.BatchNormalization()(down6)
    down7 = layers.Conv2D(512, 4, strides=2, padding='same', activation='relu')(down6)
    down7 = layers.BatchNormalization()(down7)

    # 上采样
    up1 = layers.Conv2DTranspose(512, 4, strides=2, padding='same', activation='relu')(down7)
    up1 = layers.BatchNormalization()(up1)
    up1 = layers.Concatenate()([up1, down6])

    up2 = layers.Conv2DTranspose(512, 4, strides=2, padding='same', activation='relu')(up1)
    up2 = layers.BatchNormalization()(up2)
    up2 = layers.Concatenate()([up2, down5])

    up3 = layers.Conv2DTranspose(512, 4, strides=2, padding='same', activation='relu')(up2)
    up3 = layers.BatchNormalization()(up3)
    up3 = layers.Concatenate()([up3, down4])

    up4 = layers.Conv2DTranspose(512, 4, strides=2, padding='same', activation='relu')(up3)
    up4 = layers.BatchNormalization()(up4)
    up4 = layers.Concatenate()([up4, down3])

    up5 = layers.Conv2DTranspose(256, 4, strides=2, padding='same', activation='relu')(up4)
    up5 = layers.BatchNormalization()(up5)
    up5 = layers.Concatenate()([up5, down2])

    up6 = layers.Conv2DTranspose(128, 4, strides=2, padding='same', activation='relu')(up5)
    up6 = layers.BatchNormalization()(up6)
    up6 = layers.Concatenate()([up6, down1])

    up7 = layers.Conv2DTranspose(64, 4, strides=2, padding='same', activation='relu')(up6)
    up7 = layers.BatchNormalization()(up7)

    initializer = tf.random_normal_initializer(0., 0.02)
    last = layers.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=1, padding='same',
                                  kernel_initializer=initializer, activation='tanh')(up7)

    return tf.keras.Model(inputs=inputs, outputs=last)


# 判别器模型 (PatchGAN)
def build_discriminator():
    initializer = tf.random_normal_initializer(0., 0.02)

    inp = tf.keras.layers.Input(shape=[IMG_HEIGHT, IMG_WIDTH, 3], name='input_image')

    x = layers.Conv2D(64, 4, strides=2, padding='same', kernel_initializer=initializer)(inp)
    x = layers.LeakyReLU()(x)

    x = layers.Conv2D(128, 4, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.LeakyReLU()(x)

    x = layers.Conv2D(256, 4, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.LeakyReLU()(x)

    x = layers.Conv2D(512, 4, strides=1, padding='same', kernel_initializer=initializer, use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.LeakyReLU()(x)

    last = layers.Conv2D(1, 4, strides=1, padding='same', kernel_initializer=initializer)(x)

    return tf.keras.Model(inputs=inp, outputs=last)


# 构建模型
generator_g = build_generator()  # 橙子 -> 橘子
generator_f = build_generator()  # 橘子 -> 橙子
discriminator_x = build_discriminator()  # 判别橙子
discriminator_y = build_discriminator()  # 判别橘子

# 优化器
generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

# 损失函数
loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True)


def discriminator_loss(real, generated):
    real_loss = loss_obj(tf.ones_like(real), real)
    generated_loss = loss_obj(tf.zeros_like(generated), generated)
    total_disc_loss = real_loss + generated_loss
    return total_disc_loss * 0.5


def generator_loss(generated):
    return loss_obj(tf.ones_like(generated), generated)


def calc_cycle_loss(real_image, cycled_image):
    loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))
    return LAMBDA * loss1


def identity_loss(real_image, same_image):
    loss = tf.reduce_mean(tf.abs(real_image - same_image))
    return LAMBDA * 0.5 * loss


# 由于实际数据加载复杂，这里创建模拟数据
def generate_sample_images():
    # 创建模拟的橙子和橘子图像
    orange_sample = tf.random.uniform([1, IMG_HEIGHT, IMG_WIDTH, 3], minval=0.8, maxval=1.0)  # 橙色调
    tangerine_sample = tf.random.uniform([1, IMG_HEIGHT, IMG_WIDTH, 3], minval=0.9, maxval=1.0)  # 更亮的橙色

    # 添加一些纹理
    orange_sample = orange_sample * tf.random.normal([1, IMG_HEIGHT, IMG_WIDTH, 3], mean=1.0, stddev=0.1)
    tangerine_sample = tangerine_sample * tf.random.normal([1, IMG_HEIGHT, IMG_WIDTH, 3], mean=1.0, stddev=0.05)

    return tf.clip_by_value(orange_sample, -1, 1), tf.clip_by_value(tangerine_sample, -1, 1)


# 训练循环
for epoch in range(EPOCHS):
    start = time.time()

    # 使用模拟数据
    real_x, real_y = generate_sample_images()

    with tf.GradientTape(persistent=True) as tape:
        # 生成器G将X转换为Y (橙子->橘子)
        fake_y = generator_g(real_x, training=True)
        # 生成器F将Y转换为X (橘子->橙子)
        fake_x = generator_f(real_y, training=True)

        # 循环一致性
        cycled_x = generator_f(fake_y, training=True)
        cycled_y = generator_g(fake_x, training=True)

        # 身份映射
        same_x = generator_f(real_x, training=True)
        same_y = generator_g(real_y, training=True)

        # 判别器输出
        disc_real_x = discriminator_x(real_x, training=True)
        disc_real_y = discriminator_y(real_y, training=True)
        disc_fake_x = discriminator_x(fake_x, training=True)
        disc_fake_y = discriminator_y(fake_y, training=True)

        # 计算损失
        gen_g_loss = generator_loss(disc_fake_y)
        gen_f_loss = generator_loss(disc_fake_x)

        total_cycle_loss = calc_cycle_loss(real_x, cycled_x) + calc_cycle_loss(real_y, cycled_y)

        # 总生成器损失
        total_gen_g_loss = gen_g_loss + total_cycle_loss + identity_loss(real_y, same_y)
        total_gen_f_loss = gen_f_loss + total_cycle_loss + identity_loss(real_x, same_x)

        # 判别器损失
        disc_x_loss = discriminator_loss(disc_real_x, disc_fake_x)
        disc_y_loss = discriminator_loss(disc_real_y, disc_fake_y)

    # 计算梯度
    generator_g_gradients = tape.gradient(total_gen_g_loss, generator_g.trainable_variables)
    generator_f_gradients = tape.gradient(total_gen_f_loss, generator_f.trainable_variables)
    discriminator_x_gradients = tape.gradient(disc_x_loss, discriminator_x.trainable_variables)
    discriminator_y_gradients = tape.gradient(disc_y_loss, discriminator_y.trainable_variables)

    # 应用梯度
    generator_g_optimizer.apply_gradients(zip(generator_g_gradients, generator_g.trainable_variables))
    generator_f_optimizer.apply_gradients(zip(generator_f_gradients, generator_f.trainable_variables))
    discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, discriminator_x.trainable_variables))
    discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, discriminator_y.trainable_variables))

    # 每10个epoch输出一次效果
    if (epoch + 1) % 10 == 0 or epoch == 0:
        print(f'Epoch {epoch + 1}, 生成器G损失: {total_gen_g_loss:.4f}, 生成器F损失: {total_gen_f_loss:.4f}')
        print(f'判别器X损失: {disc_x_loss:.4f}, 判别器Y损失: {disc_y_loss:.4f}')
        print(f'时间: {time.time() - start:.2f}秒')

        # 显示转换效果
        fig, axes = plt.subplots(2, 3, figsize=(12, 8))

        # 橙子 -> 橘子 转换
        axes[0, 0].imshow((real_x[0] * 0.5 + 0.5).numpy())
        axes[0, 0].set_title('原始橙子')
        axes[0, 0].axis('off')

        axes[0, 1].imshow((fake_y[0] * 0.5 + 0.5).numpy())
        axes[0, 1].set_title('转换后的橘子')
        axes[0, 1].axis('off')

        axes[0, 2].imshow((cycled_x[0] * 0.5 + 0.5).numpy())
        axes[0, 2].set_title('循环一致性')
        axes[0, 2].axis('off')

        # 橘子 -> 橙子 转换
        axes[1, 0].imshow((real_y[0] * 0.5 + 0.5).numpy())
        axes[1, 0].set_title('原始橘子')
        axes[1, 0].axis('off')

        axes[1, 1].imshow((fake_x[0] * 0.5 + 0.5).numpy())
        axes[1, 1].set_title('转换后的橙子')
        axes[1, 1].axis('off')

        axes[1, 2].imshow((cycled_y[0] * 0.5 + 0.5).numpy())
        axes[1, 2].set_title('循环一致性')
        axes[1, 2].axis('off')

        plt.suptitle(f'Epoch {epoch + 1} - CycleGAN 橙子<->橘子转换')
        plt.tight_layout()
        plt.show()

print("训练完成！")

# 最终转换效果展示
print("最终转换效果：")
test_orange, test_tangerine = generate_sample_images()
converted_tangerine = generator_g(test_orange, training=False)

fig, axes = plt.subplots(1, 2, figsize=(10, 5))
axes[0].imshow((test_orange[0] * 0.5 + 0.5).numpy())
axes[0].set_title('输入: 橙子')
axes[0].axis('off')

axes[1].imshow((converted_tangerine[0] * 0.5 + 0.5).numpy())
axes[1].set_title('输出: 转换后的橘子')
axes[1].axis('off')

plt.suptitle('CycleGAN 橙子到橘子转换 - 最终结果')
plt.tight_layout()
plt.show()